import os
import cv2
import argparse
import numpy as np
import traceback
import sys

from detectron2.engine import DefaultPredictor
from detectron2.config import get_cfg
from detectron2 import model_zoo
import torch

def extract_car(image_path, output_dir):
    try:
        print("🔧 开始配置 Detectron2...")
        # Detectron2 配置（使用官方 model_zoo 避免本地 YAML 依赖缺失）
        cfg = get_cfg()
        print("🔧 加载模型配置文件...")
        cfg.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
        cfg.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.5
        print("🔧 设置模型权重 URL...")
        cfg.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")

        # 在 Apple Silicon 优先使用 MPS，其次 CUDA，最后 CPU
        print("🔧 检查可用设备...")
        if torch.backends.mps.is_available() and torch.backends.mps.is_built():
            device = "mps"
            print("🔧 使用 MPS 设备")
        elif torch.cuda.is_available():
            device = "cuda"
            print("🔧 使用 CUDA 设备")
        else:
            device = "cpu"
            print("🔧 使用 CPU 设备")
        cfg.MODEL.DEVICE = device

        print("🔧 初始化预测器...")
        predictor = DefaultPredictor(cfg)
        print("✅ 预测器初始化完成")
    except Exception as e:
        print(f"❌ 预测器初始化失败: {e}")
        traceback.print_exc()
        sys.exit(1)

    # 读取输入图片
    img = cv2.imread(image_path)
    if img is None:
        print(f"❌ 无法读取图片: {image_path}")
        return
    outputs = predictor(img)

    # 取出 car 类别 (COCO: class 2)
    instances = outputs["instances"].to("cpu")
    car_masks = instances.pred_masks[instances.pred_classes == 2]

    if len(car_masks) == 0:
        print("❌ 没有检测到 car 对象")
        return

    # 仅选择主车（最大面积的 mask）
    areas = car_masks.flatten(1).sum(1)
    max_idx = int(areas.argmax())
    mask = car_masks[max_idx].numpy().astype(np.uint8)

    # 确保输出目录存在
    os.makedirs(output_dir, exist_ok=True)

    h, w = mask.shape

    # 转换为 RGBA
    img_rgba = cv2.cvtColor(img, cv2.COLOR_BGR2BGRA)

    # ---------- mask1.png ----------
    mask_rgba = np.zeros((h, w, 4), dtype=np.uint8)
    mask_rgba[mask == 1] = [255, 255, 255, 255]  # 车 = 白色不透明
    mask_path = os.path.join(output_dir, "mask1.png")
    cv2.imwrite(mask_path, mask_rgba)

    # ---------- front1.png ----------
    front = np.zeros_like(img_rgba)
    front[mask == 1] = img_rgba[mask == 1]
    front_path = os.path.join(output_dir, "front1.png")
    cv2.imwrite(front_path, front)

    # ---------- bg1.png ----------
    bg = img_rgba.copy()
    bg[mask == 1] = [0, 0, 0, 0]  # 车的部分透明
    bg_path = os.path.join(output_dir, "bg1.png")
    cv2.imwrite(bg_path, bg)

    print(f"✅ 已保存到 {output_dir}: bg1.png, front1.png, mask1.png")

if __name__ == "__main__":
    try:
        parser = argparse.ArgumentParser()
        parser.add_argument("-p", "--path", required=True, help="输入图片路径")
        args = parser.parse_args()

        img_path = args.path
        img_name = os.path.splitext(os.path.basename(img_path))[0]
        out_dir = os.path.join("release", img_name)

        extract_car(img_path, out_dir)
    except Exception as e:
        print(f"❌ 程序执行出错: {e}")
        traceback.print_exc()
        sys.exit(1)

    